import numpy as np
import matplotlib.pyplot as plt
import scipy.constants as sc
from functions import rstrat, PT, rho_e
from mpl_toolkits.axes_grid1 import ImageGrid
from scipy.integrate import cumtrapz 
from scipy.optimize import curve_fit

inputfile = "/home/zaire/Dropbox/Bonnie/stratification/stratification_0.61Msun_6.66Myr.dat"

r_input, P_input, T_input, rho_input, grad_ad_input, grad_rad_input, gamma_input = rstrat(inputfile)

l = 256                                              # Numer of steps in radial direction
rb = 0.15*r_input[-1]                                # Bottom of Domain
rt = 0.95*r_input[-1]                                # Top of Domain
r = np.linspace(rb , rt, l)                          # Steps used in the simulation
k = int(rb/((r[-1]-r[0])/l))                         # Just for keep r and r1 equally speced   
r1 = np.linspace(0, rb, k)                           # Steps necessary for integration 
r_t = np.hstack((r1, r))                              

# These arrays are equally spaced and varying for the bottom to the top of domain:
P = np.interp(r, r_input, P_input)
T = np.interp(r, r_input, T_input)
rho = np.interp(r, r_input, rho_input)
grad_ad = np.interp(r, r_input, grad_ad_input)
grad_rad = np.interp(r, r_input, grad_rad_input)
gamma = np.interp(r, r_input, gamma_input)       

# These are calculeted for all the extension of the radius:
P_t = np.interp(r_t, r_input, P_input)
T_t = np.interp(r_t, r_input, T_input)
rho_t = np.interp(r_t, r_input, rho_input)
grad_ad_t = np.interp(r_t, r_input, grad_ad_input)
grad_rad_t = np.interp(r_t, r_input, grad_rad_input)
gamma_t = np.interp(r_t, r_input, gamma_input)  

gamma_mean = np.mean(gamma_input[120:243])
pt = PT(T, P, gamma_mean)

# Constants
G = sc.G                                             # [G] = [N(m/kg)^2]
R = 13732.   

M = 4 * np.pi * cumtrapz(rho_t*r_t**2, r_t, initial = 0)
Msun = 1.989e+30 
#Mstar = M[-1]
Mstar = 0.8*Msun 
gb = G * M[k-1]/r[0]**2
g_t = G * M[1::] / r_t[1::]**2
#T_o, rho_o, pt_o = pol(m, r, rho[0], T[0], gb)

rhob = 3070.
Tb = 3.6e+6
# Estado Base Isentropico
m = 1.5
T_o = Tb - G * Mstar / (2 * r[-1] * R *(m + 1)) * (r**2 - r[0]**2)/ r[-1]**2  
rho_o = rhob/Tb**m * T_o**m
pt_o = T_o * (Tb * rhob / (rho_o * T_o))**(2./5.)

# Estado Ambiente Super-Adiabatico
n = 1.45
T_a = Tb - G * Mstar / (2 * r[-1] * R *(n + 1)) * (r**2 - r[0]**2)/ r[-1]**2  
rho_a = rhob/Tb**n * T_a**n
pt_a = T_a * (Tb * rhob / (rho_a * T_a))**(2./5.)

n = 1.20
T_a1 = Tb - G * Mstar / (2 * r[-1] * R *(n + 1)) * (r**2 - r[0]**2)/ r[-1]**2  
rho_a1 = rhob/Tb**n * T_a1**n
pt_a1 = T_a1 * (Tb * rhob / (rho_a1 * T_a1))**(2./5.)


for line in range(len(pt)): 
    if line == 0:
        file = open("stratification.txt", "w")
        file.write('{0:{align}{type}} {1:{align}{type}} {2:{align}{type1}} {3:{align}{type3}} {4:{align}{type1}} {5:{align}{type3}} {6:{align}{type3}}\n'.format(r[line ], T_a[line], rho_a[line], pt_a[line], T_o[line], rho_o[line], pt_o[line], align = '^', type = '0.3e', type1 = '0.2f' ,type3 = '0.8e'))
    else:
        file = open("stratification.txt", "a")
        file.write('{0:{align}{type}} {1:{align}{type}} {2:{align}{type1}} {3:{align}{type3}} {4:{align}{type}} {5:{align}{type1}} {6:{align}{type3}}\n'.format(r[line ], T_a[line], rho_a[line], pt_a[line], T_o[line], rho_o[line], pt_o[line], align = '^', type = '0.3e', type1 = '0.2f' ,type3 = '0.8e'))
        file.close()

rr = r/r_input[-1]

plt.plot(rr, T_o, 'r', rr, T_a, 'b', rr, T_a1, 'b--', rr, T, 'black')
plt.title('Temperatura')
plt.legend(('Isentropica','Ambiente n = 1.45', 'Ambiente n = 1.20','Modelo Estelar'), loc = 1)
plt.xlabel('r/Rstar')
plt.ylabel('Kelvin')
plt.savefig('T.png')
plt.show()

plt.plot(rr, rho_o, 'r', rr, rho_a, 'b', rr, rho_a1, 'b--', rr, rho, 'black')
plt.title('Densidade')
plt.legend(('Isentropica','Ambiente n = 1.45', 'Ambiente n = 1.20','Modelo Estelar'), loc = 9)
plt.xlabel('r/Rstar')
plt.ylabel('Kg/m**3')
plt.savefig('rho.png')
plt.show()


pt = pt*pt_a[0]/pt[0]

plt.plot(rr, pt_o, 'r', rr, pt_a, 'b', rr, pt_a1, 'b--', rr, pt, 'black')
plt.title('Temperatura potencial')
plt.legend(('Isentropica','Ambiente n = 1.45', 'Ambiente n = 1.20','Modelo Estelar'), loc = 6)
plt.xlabel('r/Rstar')
plt.savefig('pt.png')
plt.show()

#def pol(m, r, rhob, Tb, gb):
#    T = Tb - gb * r[0] / (2 * R * (1 + m)) * ((r/r[0])**2 - 1)
#    rho = rhob * (1 - gb * r[0]/(2 * Tb * R * (1 + m)) * ( (r/r[0])**2 - 1))**m
#    pt = T * (Tb * rhob / (rho * T))**(2./5.)
#    return T, rho, pt

#coefficients = np.polyfit(r, g_t[k-1::], 2)
#polynomial = np.poly1d(coefficients)
#g_o = polynomial(r)
